from ultralytics import YOLO
import cv2
import numpy as np
import time
import json

counter_enable=False

class VolleyballDetector:
    def __init__(self, top, left, bottom, right):
        # 加载预训练模型
        self.model = YOLO('yolov8n.pt')
        self.volleyball_count = 0
        self.last_count_time = 0
        self.count_cooldown = 1.0
        self.top = top
        self.left = left
        self.bottom = bottom
        self.right = right

        # 检测区域的阈值（相对于视频帧的百分比）
        self.detection_zone = {
            'top': self.top,      # 上边界
            'bottom': self.bottom,   # 下边界
            'left': self.left,     # 左边界
            'right': self.right     # 右边界
        }

        # 记录上一次排球的位置，用于判断方向
        self.previous_position = None

    def set_detection_zone(self, top, bottom, left, right):
        """
        设置检测区域的阈值
        :param top: 上边界百分比 (0-1)
        :param bottom: 下边界百分比 (0-1)
        :param left: 左边界百分比 (0-1)
        :param right: 右边界百分比 (0-1)
        """
        self.detection_zone = {
            'top': max(0, min(1, top)),
            'bottom': max(0, min(1, bottom)),
            'left': max(0, min(1, left)),
            'right': max(0, min(1, right))
        }

    def is_in_horizontal_zone(self, x1, x2, frame_width):
        """
        判断排球是否在水平检测区域内
        :param x1: 左边界x坐标
        :param x2: 右边界x坐标
        :param frame_width: 视频帧宽度
        :return: bool
        """
        # 计算排球中心点
        center_x = (x1 + x2) / 2

        # 计算排球中心点相对于视频帧的位置百分比
        relative_x = center_x / frame_width

        # 判断是否在水平检测区域内
        return self.detection_zone['left'] <= relative_x <= self.detection_zone['right']

    def should_count(self, y2, frame_height):
        """
        判断是否应该计数
        :param y2: 排球下边沿的y坐标
        :param frame_height: 视频帧高度
        :return: bool
        """
        if self.previous_position is None:
            self.previous_position = y2
            return False

        # 计算排球下边沿相对于视频帧的位置百分比
        relative_y = y2 / frame_height

        # 判断是否穿过下边界
        crossed_bottom = (relative_y >= self.detection_zone['bottom'] and
                         self.previous_position < self.detection_zone['bottom'])

        # 更新上一次位置
        self.previous_position = relative_y

        return crossed_bottom

    def detect_objects(self, frame):
        """
        检测单帧中的人物和排球
        :param frame: 输入帧
        :return: 检测结果字典
        """
        frame_height, frame_width = frame.shape[:2]

        # 进行检测
        results = self.model(frame)

        # 初始化返回结果
        detection_results = {
            'person': None,
            'volleyball': None,
            'count': self.volleyball_count,
            'detection_zone': {
                'top': int(self.detection_zone['top'] * frame_height),
                'bottom': int(self.detection_zone['bottom'] * frame_height),
                'left': int(self.detection_zone['left'] * frame_width),
                'right': int(self.detection_zone['right'] * frame_width)
            }
        }

        # 处理检测结果
        for result in results:
            boxes = result.boxes
            for box in boxes:
                x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
                conf = box.conf[0].cpu().numpy()
                cls = int(box.cls)

                # 检测人物（类别索引为0）
                if cls == 0:
                    detection_results['person'] = {
                        'bbox': [int(x1), int(y1), int(x2), int(y2)],
                        'confidence': float(conf)
                    }

                # 检测排球（类别索引为32）
                elif cls == 32 and conf > 0.5:
                    detection_results['volleyball'] = {
                        'bbox': [int(x1), int(y1), int(x2), int(y2)],
                        'confidence': float(conf)
                    }

                    # 判断排球是否在水平检测区域内
                    if self.is_in_horizontal_zone(x1, x2, frame_width):
                        current_time = time.time()
                        if (self.should_count(y2, frame_height) and
                            current_time - self.last_count_time > self.count_cooldown):
                            self.volleyball_count += 1
                            self.last_count_time = current_time
                            detection_results['count'] = self.volleyball_count
        
        return detection_results


def enable_counter(v):
    global enable_counter
    counter_enable = v

def process_video(source=0,websocket=None,output_path=None):
    """
    处理视频流
    :param source: 视频源（可以是摄像头索引或视频文件路径）
    """
    # 创建检测器实例
    detector = VolleyballDetector(0.1,0.2,0.3,0.7)
    
    # 打开视频源
    cap = cv2.VideoCapture(source)
    if not cap.isOpened():
        raise ValueError(f"无法打开视频源: {source}")

    # 创建视频写入器（如果指定了输出路径）
    writer = None
    if output_path:
        # 获取视频属性
        frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        fps = int(cap.get(cv2.CAP_PROP_FPS))        
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        writer = cv2.VideoWriter(output_path, fourcc, fps, (frame_width, frame_height)) 
    
    counter = 0
    try:
        while counter_enable:
            ret, frame = cap.read()
            if not ret:
                break
               
            # 进行检测并获取结果
            results = detector.detect_objects(frame)
            
            # 打印检测结果
            print("检测结果:")
            if results['person']:
                print(f"人物位置: {results['person']['bbox']}")
            if results['volleyball']:
                print(f"排球位置: {results['volleyball']['bbox']}")
            print(f"计数: {results['count']}")
            print("-" * 50)
            if websocket is not None and counter != results['count']:
                counter = results['count']
                websocket.send(json.dumps({
                                    'type': 'counter_update',
                                    'count': counter
                }))
            # 写入输出视频
            if writer:
                writer.write(frame)    
    finally:
        # 释放资源
        cap.release()
        if writer:
            writer.release()
        writer = None

def main():
    # 可以选择使用摄像头（0表示默认摄像头）或视频文件
    source = 0  # 使用摄像头
    source = "1.mp4"  # 使用视频文件
    
    try:
        process_video(source)
    except Exception as e:
        print(f"发生错误: {str(e)}")

if __name__ == "__main__":
    main() 